import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
from tftb.processing.cohen import WignerVilleDistribution, PageRepresentation, MargenauHillDistribution
from scipy.signal import hilbert, spectrogram, get_window
import numpy as np
from numpy import log as ln
def one_3d_pick(x,y,z,len_step):

    fig = plt.figure(figsize=(10, 10))
    ax = plt.axes(projection='3d')
    ax.scatter(x[:len_step], y[:len_step], z[:len_step], marker=">", c='black')

    plt.show()
def Anime(x, y, z, len_step):

    from matplotlib.animation import PillowWriter

    metadata = dict(title="Movie")
    writer = PillowWriter(fps=2, metadata=metadata)

    fig = plt.figure(figsize=(12, 12), dpi=100)
    ax = fig.add_subplot(projection='3d')

    def update(i):
        ax.clear()
        ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step], z[i * len_step:i * len_step + len_step], s=0.5)
        plt.title(str(i * len_step))
        plt.xlabel('$Bx$', color='g')
        plt.ylabel('$By$', color='g')
        ax.set_zlabel('$Bz$')

    ani = animation.FuncAnimation(fig, update, np.arange(100), interval=1000, repeat=False)

    with writer.saving(fig, "Banimate-fix.gif", 100):
        for i in range(100):
            ax.clear()
            ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step], z[i * len_step:i * len_step + len_step])
            plt.title(str(i * len_step))
            plt.xlabel('$Bx$', color='g')
            plt.ylabel('$By$', color='g')
            ax.set_zlabel('$Bz$')
            writer.grab_frame()
    plt.show()

def Video(x, y, z, len_step):

    from matplotlib.animation import FFMpegWriter

    metadata = dict(title="Movie")
    writer2 = FFMpegWriter(fps=2, metadata=metadata)

    fig = plt.figure(figsize=(12, 12), dpi=100)
    ax = fig.add_subplot(projection='3d')

    with writer2.saving(fig, "Banimate.mp4", 100):
        for i in range(100):
            ax.clear()
            ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step], z[i * len_step:i * len_step + len_step])
            plt.title(str(i*len_step))
            plt.xlabel('$Bx$', color='g')
            plt.ylabel('$By$', color='g')
            ax.set_zlabel('$Bz$')
            # ax.set_xlim(min(X), max(X))
            # ax.set_ylim(min(Y), max(Y))
            # ax.set_zlim(min(Z), max(Z))
            writer2.grab_frame()
    plt.show()

    fig = plt.figure(figsize=plt.figaspect(0.5))
    amoung = 3
    j = 1
    for i in range(24, 33):
        ax = fig.add_subplot(amoung, amoung, j, projection='3d')
        ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step], z[i * len_step:i * len_step + len_step], s=0.01)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        ax.set_zticklabels([])
        j += 1
    plt.show()

def Diagrams():
    X = df["Bx_GSE"]
    Y = df["By_GSE"]
    Z = df["Bz_GSE"]

    step = 5000

    one_3d_pick(X,Y,Z,step)


    #Anime(X, Y, Z, step)


def ReadTable(file):
    arr = ["Year", "Day", "Hour", "Minute", "Sec", "Bx_GSE", "By_GSE", "Bz_GSE"]
    df = pd.read_table(file, names=arr, sep="\s+")
    for i in arr:
        df.loc[(df[i] > 9999)] = 0
    return df

df = ReadTable('sun.lst')
Diagrams()